// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <melon/benchmark.h>
#include <melon/init/init.h>

#include <pollux/benchmarks/ExpressionBenchmarkBuilder.h>
#include <pollux/functions/lib/benchmarks/FunctionBenchmarkBase.h>
#include <pollux/functions/prestosql/registration/registration_functions.h>

using namespace kumo::pollux;
using namespace kumo::pollux::exec;
using namespace kumo::pollux::functions;

int main(int argc, char** argv) {
  melon::Init init(&argc, &argv);

  functions::prestosql::registerArrayFunctions();
  memory::MemoryManager::initialize({});
  auto anotherPool = memory::memoryManager()->addLeafPool("bm");

  ExpressionBenchmarkBuilder benchmarkBuilder;

  auto* pool = benchmarkBuilder.pool();
  auto& vm = benchmarkBuilder.vectorMaker();

  auto createSet = [&](const TypePtr& type,
                       bool withNulls,
                       const VectorPtr& constantInput,
                       const std::string& suffix = "") {
    VectorFuzzer::Options options;
    options.vectorSize = 1'000;
    options.nullRatio = withNulls ? 0.2 : 0.0;

    VectorFuzzer fuzzer(options, pool);
    std::vector<VectorPtr> columns;
    columns.push_back(fuzzer.fuzzFlat(type));
    columns.push_back(fuzzer.fuzzFlat(type));
    columns.push_back(fuzzer.fuzzFlat(type));
    columns.push_back(
        BaseVector::create_null_constant(type, options.vectorSize, pool));
    columns.push_back(
        BaseVector::wrap_in_constant(options.vectorSize, 0, constantInput));

    auto input = vm.row_vector({"c0", "c1", "c2", "n", "c"}, columns);

    benchmarkBuilder
        .addBenchmarkSet(
            fmt::format(
                "array_constructor_{}_{}{}",
                mapTypeKindToName(type->kind()),
                withNulls ? "nulls" : "nullfree",
                suffix),
            input)
        .addExpression("1", "array_constructor(c0)")
        .addExpression("2", "array_constructor(c0, c1)")
        .addExpression("3", "array_constructor(c0, c1, c2)")
        .addExpression("2_null", "array_constructor(c0, c1, n)")
        .addExpression("2_const", "array_constructor(c0, c1, c)");
  };

  auto constantInteger = BaseVector::create_constant(INTEGER(), 11, 1, pool);
  createSet(INTEGER(), true, constantInteger);
  createSet(INTEGER(), false, constantInteger);

  {
    auto constantString =
        BaseVector::create_constant(VARCHAR(), std::string(20, 'x'), 1, pool);
    createSet(VARCHAR(), true, constantString);
    createSet(VARCHAR(), false, constantString);
  }

  pool = anotherPool.get();
  {
    auto constantString =
        BaseVector::create_constant(VARCHAR(), std::string(20, 'x'), 1, pool);
    createSet(VARCHAR(), true, constantString, "_dp");
    createSet(VARCHAR(), false, constantString, "_dp");
  }

  pool = benchmarkBuilder.pool();

  auto constant_row = vm.row_vector({
      BaseVector::create_constant(INTEGER(), 11, 1, pool),
      BaseVector::create_constant(DOUBLE(), 1.23, 1, pool),
  });
  createSet(ROW({INTEGER(), DOUBLE()}), true, constant_row);
  createSet(ROW({INTEGER(), DOUBLE()}), false, constant_row);

  auto constantArray = vm.array_vector<int32_t>({{1, 2, 3, 4, 5}});
  createSet(ARRAY(INTEGER()), true, constantArray);
  createSet(ARRAY(INTEGER()), false, constantArray);

  auto constantMap = vm.map_vector<int32_t, float>({{{1, 1.23}, {2, 2.34}}});
  createSet(MAP(INTEGER(), REAL()), true, constantMap);
  createSet(MAP(INTEGER(), REAL()), false, constantMap);

  benchmarkBuilder.registerBenchmarks();

  melon::runBenchmarks();
  return 0;
}
